import numpy as np
from numpy import linalg as LA
from numpy.linalg import inv
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from loss import loss_be
import torch

def FF(U):
    return U * (1 - U * U)
    # return 0


def circle_x(P, U):
    return np.matmul(P, U)


def circle_y(P, U):
    return np.matmul(U, P)

def mass_engy(U, hx, hy,epsilon):
    mass = 0
    egy1 = 0
    egy2 = 0
    N=np.shape(U)[0]
    M=np.shape(U)[1]
    for i in range(N - 1):
        for j in range(M - 1):
            uu = 0.25 * (U[i, j] + U[i + 1, j] + U[i, j + 1] + U[i + 1, j + 1])
            mass = mass + uu
            egy1 = egy1 + (uu * uu - 1) ** 2
            udx = 0.5 * (U[i + 1, j] - U[i, j] + U[i + 1, j + 1] - U[i, j + 1]) / hx
            udy = 0.5 * (U[i, j + 1] - U[i, j] + U[i + 1, j + 1] - U[i + 1, j]) / hy
            egy2 = egy2 + (udx * udx + udy * udy)

    mass = mass * hx * hy
    egy1 = 0.25 * egy1
    egy2 = 0.5 * egy2 * epsilon*epsilon
    egy = (egy1 + egy2) * hx * hy
    return mass, egy

def U_full(U):
    N = np.shape(U)[0]
    M = np.shape(U)[1]
    U_full = np.zeros((N + 1, M + 1))
    U_full[1:N + 1, :M] = U[:N, :M]
    U_full[0, :M] = U[N-1, :M]
    U_full[:, M] = U_full[:, 0]
    return U_full


def numerical(U,domain,n,dt,epsilon,type='BE'):
    A = U.shape
    N = A[0]
    M = A[1]

    lx = domain[1] - domain[0]
    ly = domain[3] - domain[2]

    Vol = lx * ly
    hx = lx / N
    hy = ly / M
    K = n
    A = np.zeros((N,N))
    B = np.zeros((M,M))
    for i in range(N - 1):
        A[i, i] = -2
        A[i, i + 1] = 1
        A[i + 1, i] = 1
    A[N - 1, N - 1] = -2
    A[0, N - 1] = 1
    A[N - 1, 0] = 1

    for j in range(M - 1):
        B[j, j] = -2
        B[j, j + 1] = 1
        B[j + 1, j] = 1
    B[M - 1, M - 1] = -2
    B[0, M - 1] = 1
    B[M - 1, 0] = 1


    A = epsilon ** 2 * A/(hx*hx)
    [Dx,Px] = LA.eig(A)
    Pxi = inv(Px)
    # a= np.matmul(np.matmul(Px,np.diag(Dx)),Pxi)
    B = epsilon ** 2 * B / (hy * hy)
    [Dy,Py] = LA.eig(B)
    Pyi = inv(Py)


    iter_max=50


    if type == 'BE':
        DD = np.zeros((N, M))

        for i in range(N):
            for j in range(M):
                DD[i, j] = 1 - dt* (Dx[i] + Dy[j])
        a, b = mass_engy(U_full(U), hx, hy,epsilon)
        mass = [a]
        egy = [b]

        for k in range(n):
            RH_curr = U

            # u_n = torch.unsqueeze(torch.unsqueeze(torch.from_numpy(U), 0), 0)
            for iter in range(iter_max):

                FU=FF(U)
                RH = RH_curr+dt*FU
                U_new = circle_y(Pyi,circle_x(Px,np.divide(circle_y(Py,circle_x(Pxi,RH)),DD)))

                if (np.max(np.abs(U-U_new))) < 1e-8:
                    U= U_new
                    break
                U=U_new

            # u_next = torch.unsqueeze(torch.unsqueeze(torch.from_numpy(U_new), 0), 0)
            # L = loss_be(u_n, u_next, dt, domain, epsilon, type='BE')[0][0].cpu().detach().numpy()
            #
            # P = np.linalg.norm(L)
            # print(P)

            U_full1 = U_full(U)
            a,b = mass_engy(U_full1,hx,hy,epsilon)
            mass.append(a)
            egy.append(b)
            # if (np.abs(egy[-1]-egy[-2])<max(egy[-2],1)*1e-6):
            #     n=k
            #     break


    if type =='CN':
        DD = np.zeros((N, M))
        for i in range(N):
            for j in range(M):
                DD[i, j] = 1 - 0.5*dt * (Dx[i] + Dy[j])

        a, b = mass_engy(U_full(U), hx, hy,epsilon)
        mass = [a]
        egy = [b]
        for k in range(int(n)):
            FU=FF(U)
            RH_curr = U+0.5*dt*(circle_x(A,U)+circle_y(B,U))+0.5*dt*FU

            # u_n = torch.unsqueeze(torch.unsqueeze(torch.from_numpy(U), 0), 0)

            for iter in range(iter_max):
                FU=FF(U)
                RH = RH_curr+0.5*dt*FU
                U_new = circle_y(Pyi,circle_x(Px,np.divide(circle_y(Py,circle_x(Pxi,RH)),DD)))
                if (np.max(abs(U-U_new))) < 1e-8:
                    U= U_new
                    break
                U=U_new

            # u_next = torch.unsqueeze(torch.unsqueeze(torch.from_numpy(U_new), 0), 0)
            # L = loss_be(u_n, u_next, dt, domain, epsilon, type='CN')[0][0].cpu().detach().numpy()
            # P=np.linalg.norm(L)
            # P1=np.max(abs(L))
            # print(P,P1)

            U_full1 = U_full(U)
            a, b = mass_engy(U_full1, hx, hy,epsilon)
            mass.append(a)
            egy.append(b)
            # if (np.abs(egy[-1] - egy[-2]) < max(egy[-2], 1) * 1e-6):
            #     n = k
            #     break
    plt.plot(np.linspace(1, len(egy), len(egy)), egy,label ='BE energy')
    plt.legend()
    plt.xlabel('Timesteps')
    plt.ylabel('energy')
    plt.show()

    return U_full1


# H=W=128
# dt=0.001
# domain=[0,1,0,1]
# epsilon = 0.01
# type = 'CN'
# n=10000
# # #
# xx = np.linspace(0, 1, H+1)
# yy = np.linspace(0, 1, W+1)
#
# X, Y = np.meshgrid(xx, yy)
# Z = np.tanh((0.35 - np.sqrt((X-0.5)*(X-0.5) + (Y-0.5) * (Y-0.5))) / (2*epsilon))
# A=np.zeros((H,W))
# A=Z[1:H+1,:W]
#
#
# b=numerical(A, domain, n,dt,epsilon, type=type)
# np.save(f'data/deltaT{dt}_{H}_{W}_{n}_{type}_full.npy', b)



# H=W=128
# dt=0.5
# domain=[0,1,0,1]
# epsilon = 0.01
# n=100
# #
# xx = np.linspace(0, 1, H+1)
# yy = np.linspace(0, 1, W+1)
#
# X, Y = np.meshgrid(xx, yy)
# Z = np.tanh((0.35 - np.sqrt((X-0.5)*(X-0.5) + (Y-0.5) * (Y-0.5))) / 0.02)
# A=np.zeros((H,W))
# A=Z[1:H+1,:W]
# b=numerical(A, domain, n,dt,epsilon, type='CN')
# # # # # Zexact = np.exp(-5*np.pi**2*dt*n)*np.sin(np.pi*(X-1/4))*np.sin(2*np.pi*(Y-1/8))
# # # # # Z_exact = U_full(Zexact)
# # # # fig, ax = plt.subplots()
# # # # im = ax.imshow(Z, interpolation='bilinear', cmap=cm.RdYlGn)
# # # # plt.title('Z')
# # # # plt.colorbar(im)
# # # #
# fig, ax1 = plt.subplots()
# im1 = ax1.imshow(b, interpolation='bilinear', cmap=cm.RdYlGn)
# plt.title('b')
# plt.colorbar(im1)
# #
# plt.show()



# np.random.seed(seed=10)
# U=(-1+np.random.rand(128,128)*2)*0.9
# domain = [-0.5,0.5,-0.5,0.5]
#
# epsilon =0.01
# # # a=numerical(U, domain, 10, deltaT, epsilon, type='CN')
# b=numerical(U, [-0.5,0.5,-0.5,0.5], 100,0.001,0.01, type='CN')
# x=[1,2,4,8,16,32]
#
# for i in range(6):
#         ai = numerical(U, [-0.5,0.5,-0.5,0.5], x[i], 0.1/x[i], 0.01, type='BE')
#         ci = np.linalg.norm(ai-b,'fro')/np.linalg.norm(b,'fro')
#         print(ci)